import ftplib
import os
import concurrent.futures
import time

# --- FTP服务器连接信息 ---
FTP_HOST = "ftp.ptree.jaxa.jp"
FTP_USER = "********"
FTP_PASS = "****"

# --- 用户配置区域 ---
REMOTE_BASE_PATH = "/jma/netcdf/202506/24"
LOCAL_BASE_PATH = r"E:\zwh_data\Fire_nc_data"

# 下载指定后缀的文件
FILES_TO_DOWNLOAD = ["06001.nc"]

# 文件所在目录深度
TARGET_FILE_PARENT_DEPTH = 0

# 最大并发下载数（视情况而定，避免服务器拒绝）
MAX_WORKERS = 16

# 连接超时设置（秒）
FTP_TIMEOUT = 30

# 单个文件下载重试次数
MAX_RETRIES = 3


def connect_ftp(host, user, passwd, timeout=FTP_TIMEOUT):
    """连接到FTP服务器并切换为二进制模式"""
    try:
        print(f"尝试连接到FTP服务器: {host}...")
        ftp = ftplib.FTP(host, timeout=timeout)
        ftp.login(user, passwd)

        # 切换为二进制传输模式（解决SIZE命令错误）
        ftp.voidcmd('TYPE I')  # 切换到二进制模式（TYPE I）
        print(f"成功连接到FTP服务器并切换为二进制模式: {host}")
        return ftp
    except ftplib.all_errors as e:
        print(f"连接FTP服务器失败: {e}")
        return None


def download_file(ftp_host, ftp_user, ftp_pass, remote_file_dir, remote_filename, local_filepath):
    """从FTP下载单个文件（二进制模式）"""
    full_remote_path = f"{remote_file_dir.rstrip('/')}/{remote_filename}"
    print(f"[下载任务] 开始处理: {full_remote_path}")

    for attempt in range(1, MAX_RETRIES + 1):
        ftp_conn_dl = None
        try:
            # 检查本地目录
            local_dir = os.path.dirname(local_filepath)
            if not os.path.exists(local_dir):
                os.makedirs(local_dir)
                print(f"  创建本地目录: {local_dir}")

            # 建立连接（已默认二进制模式）
            ftp_conn_dl = connect_ftp(ftp_host, ftp_user, ftp_pass)
            if not ftp_conn_dl:
                raise Exception("无法建立FTP连接")

            # 切换到远程目录
            ftp_conn_dl.cwd(remote_file_dir)

            # 获取远程文件大小
            try:
                remote_size = ftp_conn_dl.size(remote_filename)
                print(f"  远程文件大小: {remote_size} 字节")
            except ftplib.error_perm:
                # 如果服务器仍不支持SIZE命令，跳过大小检查
                print(f"  服务器不支持SIZE命令，跳过大小检查")
                remote_size = None

            # 开始下载（二进制模式传输）
            print(f"  尝试下载 (尝试 {attempt}/{MAX_RETRIES}): {full_remote_path}")
            start_time = time.time()

            with open(local_filepath, 'wb') as local_file:
                # 使用二进制传输命令RETR
                ftp_conn_dl.retrbinary(f"RETR {remote_filename}", local_file.write)

            # 验证下载完整性（如果获取到远程大小）
            if remote_size:
                local_size = os.path.getsize(local_filepath)
                if local_size != remote_size:
                    raise Exception(f"下载不完整（本地: {local_size} 字节，远程: {remote_size} 字节）")

            # 下载成功
            duration = time.time() - start_time
            print(f"成功下载 {full_remote_path} 到 {local_filepath}")
            print(f"  耗时: {duration:.2f} 秒")
            return True

        except ftplib.all_errors as e:
            print(f"下载文件失败 (尝试 {attempt}/{MAX_RETRIES}): {e}")
            # 清理不完整文件
            if os.path.exists(local_filepath):
                try:
                    os.remove(local_filepath)
                    print(f"  已删除不完整文件: {local_filepath}")
                except OSError as oe:
                    print(f"清理部分下载文件失败: {oe}")
            # 重试等待
            if attempt < MAX_RETRIES:
                wait_time = 2 ** attempt
                print(f"  将在 {wait_time} 秒后重试...")
                time.sleep(wait_time)
        finally:
            if ftp_conn_dl:
                try:
                    ftp_conn_dl.quit()
                except ftplib.all_errors:
                    pass

    print(f"下载文件 {full_remote_path} 失败，已达到最大重试次数")
    return False


def _process_ftp_directory_recursively(ftp, current_level_local_dir, files_spec, current_depth, file_container_depth,
                                       executor, ftp_host, ftp_user, ftp_pass):
    """递归处理FTP目录，提交下载任务"""
    print(f"[深度 {current_depth}] 进入远程目录: {ftp.pwd()}")
    print(f"[深度 {current_depth}] 对应的本地目录: {current_level_local_dir}")

    os.makedirs(current_level_local_dir, exist_ok=True)

    try:
        items = ftp.nlst()
        print(f"[深度 {current_depth}] 在远程目录中找到 {len(items)} 个项目")

        for item_name in items:
            if item_name in ['.', '..']:
                continue

            # 判断是文件还是目录
            is_dir = False
            try:
                original_dir = ftp.pwd()
                ftp.cwd(item_name)
                is_dir = True
                ftp.cwd(original_dir)
            except ftplib.error_perm as e_perm:
                if "550" in str(e_perm):
                    is_dir = False
                else:
                    print(f"[深度 {current_depth}] 检查项目 '{item_name}' 类型时出错: {e_perm}")
                    continue

            if is_dir:
                print(f"[深度 {current_depth}] 发现子目录: {item_name}")
                if current_depth < file_container_depth:
                    _process_ftp_directory_recursively(
                        ftp,
                        os.path.join(current_level_local_dir, item_name),
                        files_spec,
                        current_depth + 1,
                        file_container_depth,
                        executor,
                        ftp_host,
                        ftp_user,
                        ftp_pass
                    )
            else:
                if current_depth == file_container_depth and any(item_name.endswith(ext) for ext in files_spec):
                    print(f"[深度 {current_depth}] 找到目标文件: {item_name}")
                    local_filepath = os.path.join(current_level_local_dir, item_name)
                    executor.submit(download_file, ftp_host, ftp_user, ftp_pass, ftp.pwd(), item_name, local_filepath)
                else:
                    print(f"[深度 {current_depth}] 跳过不匹配的文件: {item_name}")

    except ftplib.all_errors as e:
        print(f"[深度 {current_depth}] 列出远程目录内容失败: {e}")


def main():
    # 创建本地目录
    os.makedirs(LOCAL_BASE_PATH, exist_ok=True)

    # 连接FTP并切换到二进制模式
    ftp_connection = connect_ftp(FTP_HOST, FTP_USER, FTP_PASS)
    if ftp_connection:
        try:
            ftp_connection.cwd(REMOTE_BASE_PATH)
            print(f"成功切换到远程基础路径: {ftp_connection.pwd()}")

            # 统计目标文件
            items = ftp_connection.nlst()
            target_files = [item for item in items if any(item.endswith(ext) for ext in FILES_TO_DOWNLOAD)]
            print(f"找到 {len(target_files)} 个目标文件，准备下载...")

            # 启动下载
            with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
                _process_ftp_directory_recursively(
                    ftp_connection,
                    LOCAL_BASE_PATH,
                    FILES_TO_DOWNLOAD,
                    0,
                    TARGET_FILE_PARENT_DEPTH,
                    executor,
                    FTP_HOST,
                    FTP_USER,
                    FTP_PASS
                )
                print("所有下载任务已提交，等待完成...")

            print("所有下载任务处理完毕")

        except ftplib.all_errors as e:
            print(f"主流程FTP错误: {e}")
        # finally:
        #     ftp_connection.quit()
        #     print("已断开FTP连接")
    else:
        print("无法连接FTP服务器，退出")


if __name__ == "__main__":
    main()